#include <iostream>
#include <cassert>
#include <string.h>
#include <unistd.h>
#include <vector>
#include <fstream>
#include <sstream>
#include "model_crypto.h"

std::string load_model(const std::string &model_path)
{
  // deserialize the .engine and run inference
  std::ifstream file(model_path, std::ios::binary);
  if (!file.good())
  {
    std::cerr << "read " << model_path << " error!" << std::endl;
  }
  std::ostringstream tmp;
  tmp << file.rdbuf();
  return tmp.str();
}

int save_model(const std::string &model_path, const std::string &model_context)
{
  std::ofstream file(model_path, std::ios::binary);
  if (!file)
  {
    std::cerr << "could not open plan output file" << std::endl;
    return -1;
  }
  file.write(reinterpret_cast<const char *>(model_context.data()),
             model_context.length());
  file.close();
}

int main(int argc, char **argv)
{
  if (argc < 2)
  {
    std::cerr << "please input model path" << std::endl;
    return false;
  }
  std::string model_path = argv[1];
  std::string model_context;
  model_context = load_model(model_path);

  auto encrypt_model_context = encrypt_model(model_context);
  save_model("encrypt_model.trt", encrypt_model_context);

  auto decrypt_model_context = decrypt_model(encrypt_model_context);
  save_model("decrypt_model.trt", decrypt_model_context);

  if (model_context == encrypt_model_context)
    std::cout << "false" << std::endl;

  if (model_context == decrypt_model_context)
    std::cout << "success" << std::endl;
  return 0;
}